package com.iat.demo.client;

import com.google.protobuf.ByteString;
import com.iat.demo.proto.protocol.IatGrpc;
import com.iat.demo.proto.protocol.IatRequest;
import com.iat.demo.proto.protocol.IatResult;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * IAT语音识别客户端
 */
public class IatClient {

    private static final String IAT_HOST = "42.62.116.141";
    private static final int IAT_PORT = 9097;
    private static final int CHUNK_SIZE = 1024 * 10; // 每个音频分片大小为10KB

    public static void main(String[] args) {
        // 默认使用audio/hftq.wav文件
        String audioFilePath = "audio/20s.pcm";
        
        // 如果提供了命令行参数，则使用命令行参数指定的文件
        if (args.length > 0) {
            audioFilePath = args[0];
        }

        try {
            System.out.println("开始识别音频文件: " + audioFilePath);
            new IatClient().recognizeAudio(audioFilePath);
        } catch (Exception e) {
            System.err.println("识别过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 识别音频文件
     *
     * @param audioFilePath 音频文件路径
     * @throws Exception 识别异常
     */
    public void recognizeAudio(String audioFilePath) throws Exception {
        // 创建gRPC通道
        ManagedChannel channel = ManagedChannelBuilder.forAddress(IAT_HOST, IAT_PORT)
                .usePlaintext()
                .build();

        try {
            // 创建异步客户端
            IatGrpc.IatStub asyncStub = IatGrpc.newStub(channel);

            // 创建一个CountDownLatch来等待识别完成
            CountDownLatch finishLatch = new CountDownLatch(1);

            // 创建响应观察者
            StreamObserver<IatResult> responseObserver = createResponseObserver(finishLatch);

            // 创建请求观察者
            StreamObserver<IatRequest> requestObserver = asyncStub.createRec(responseObserver);

            // 发送音频数据
            sendAudioData(audioFilePath, requestObserver);

            // 等待识别完成，最多等待2分钟
            if (!finishLatch.await(2, TimeUnit.MINUTES)) {
                System.err.println("识别超时，未在预期时间内完成");
            }
        } finally {
            // 关闭通道
            channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
        }
    }

    /**
     * 创建响应观察者
     *
     * @param finishLatch 完成标记
     * @return 响应观察者
     */
    private StreamObserver<IatResult> createResponseObserver(CountDownLatch finishLatch) {
        return new StreamObserver<IatResult>() {
            @Override
            public void onNext(IatResult result) {
                if (result.getErrCode() == 0) {
                    System.out.println("识别结果: " + result.getAnsStr());
                } else {
                    System.err.println("识别错误: [" + result.getErrCode() + "] " + result.getErrStr());
                }

                if (result.getEndFlag()) {
                    System.out.println("识别结束");
                    finishLatch.countDown();
                }
            }

            @Override
            public void onError(Throwable t) {
                System.err.println("识别过程中发生错误: " + t.getMessage());
                finishLatch.countDown();
            }

            @Override
            public void onCompleted() {
                System.out.println("识别完成");
                finishLatch.countDown();
            }
        };
    }

    /**
     * 发送音频数据
     *
     * @param audioFilePath    音频文件路径
     * @param requestObserver  请求观察者
     * @throws IOException 读取音频文件异常
     */
    private void sendAudioData(String audioFilePath, StreamObserver<IatRequest> requestObserver) throws IOException {
        // 读取音频文件，分片处理
        byte[][] audioChunks = AudioUtils.readAudioFileByChunk(audioFilePath, CHUNK_SIZE);
        System.out.println("音频文件分片完成，共 " + audioChunks.length + " 个分片");

        // 发送第一个分片，包含会话参数
        IatRequest firstRequest = createFirstRequest(audioChunks[0]);
        requestObserver.onNext(firstRequest);
        System.out.println("发送第一个音频分片，包含会话参数");

        // 发送其余分片
        for (int i = 1; i < audioChunks.length; i++) {
            IatRequest request = IatRequest.newBuilder()
                    .setSamples(ByteString.copyFrom(audioChunks[i]))
                    .setEndFlag(i == audioChunks.length - 1)
                    .build();

            requestObserver.onNext(request);
            System.out.println("发送音频分片 " + (i + 1) + "/" + audioChunks.length +
                    (i == audioChunks.length - 1 ? " (最后一个分片)" : ""));

            // 稍作延迟，避免发送过快
            try {
                Thread.sleep(10);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

        // 标记请求完成
        requestObserver.onCompleted();
    }

    /**
     * 创建第一个请求，包含会话参数
     *
     * @param audioData 音频数据
     * @return 请求对象
     */
    private IatRequest createFirstRequest(byte[] audioData) {
        Map<String, String> sessionParams = new HashMap<>();
        sessionParams.put("sid", UUID.randomUUID().toString());
        sessionParams.put("eos", "60000");
        sessionParams.put("aue", "raw");
        sessionParams.put("rst", "plain");

        return IatRequest.newBuilder()
                .putAllSessionParam(sessionParams)
                .setSamples(ByteString.copyFrom(audioData))
                .setEndFlag(false)
                .build();
    }
} 